package tbdp.tool.writer;


import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataOutputStream;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.IOUtils;

import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;

import tbdp.tool.common.JobConstants;
import tbdp.tool.config.WriterConfig;
import tbdp.tool.container.WriterContainer;
import tbdp.tool.context.JobContext;
import tbdp.tool.spi.RecordIterator;
import tbdp.tool.spi.Writer;

/**
 * @author xbkaishui
 * @version $Id: HdfsWriter.java, v 0.1 2017-06-18 下午12:20 xbkaishui Exp $$
 */
public class HdfsWriter extends Writer {

    private Configuration conf = new Configuration();

    private Writer delegateWriter;

    @Override
    public void write(JobContext context, RecordIterator recordIterator) {
        String format = getWriterFormat(context);
        delegateWriter = WriterContainer.getWriter(format);
        delegateWriter.write(context, recordIterator);
    }

    @Override
    public void postWriter(JobContext context, RecordIterator recordIterator) {
        String tmpFilePath = getTmpFilePath(context);
        String localFilePath = getFilePath(context);
        //sync to hdfs
        saveFileToHdfs(tmpFilePath, localFilePath);
    }

    /**
     * 文件保存到hdfs
     */
    protected void saveFileToHdfs(String tmpFilePath, String output) {
        logger.info("save local file {}   to hdfs {} ", tmpFilePath, output);
        FSDataOutputStream outputStream = null;
        InputStream in = null;
        try {
            FileSystem fs = FileSystem.get(conf);
            outputStream = fs.create(new Path(output));
            in = new FileInputStream(tmpFilePath);
            IOUtils.copyBytes(in, outputStream, conf);

        } catch (IOException e) {
            throw new RuntimeException(e);
        } finally {
            IOUtils.closeStream(in);
            IOUtils.closeStream(outputStream);
        }
    }

    /**
     * 获取 writer 格式
     */
    private String getWriterFormat(JobContext context) {
        WriterConfig writer = context.getJobConfig().getWriter();
        return writer.getConfigPropAsString(JobConstants.FORMAT, "csv");
    }

    @Override
    public void close(JobContext context) {
        if (delegateWriter != null) {
            delegateWriter.close(context);
        }
    }

    @Override
    public String getType() {
        return "hdfs";
    }
}
